from io import BytesIO

import numpy as np
import torch
from PIL import Image, ImageFilter
from skimage.metrics import structural_similarity as SSIM

from util.seed import set_seed
from raw.get_model import get_model
from attack.get_attack import get_attack
from util.image_convert import np2pil, pil2np

class AttackEvaluation:
    def __init__(self, model, attack, suffix = ''):
        self.model = model
        self.attack = attack

        # 加载训练好的网络
        model.load_raw_state()

        # 加载攻击前目标数据
        base_dir = 'data/candidate_data/' + model.model_name + '/'
        self.images = np.load(base_dir + 'images.npy')
        self.labels = np.load(base_dir + 'labels.npy')
        if attack.targeted:
            if attack.llc:
                self.targets = np.load(base_dir + 'llcs.npy')
            else:
                self.targets = np.load(base_dir + 'targets.npy')
        else:
            self.targets = None

        # 加载攻击后目标数据
        base_dir = 'data/adversarial_data/' + model.model_name + '_' + attack.name + suffix + '/'
        self.adv_images = np.load(base_dir + 'images.npy')

        # 用攻击后图像作为输入，计算网络输出的softmax
        outs = model.predict(self.adv_images)
        exps = np.exp(outs)
        arr = []
        for i in range(len(outs)):
            arr.append(exps[i] / np.sum(exps[i]))
        self.softmax_outs = np.array(arr)

    # 某个攻击后图像是否成功误导
    def successful_adv(self, index):
        pred = np.argmax(self.softmax_outs[index])
        if self.attack.targeted:
            return pred == self.targets[index]
        else:
            return pred != self.labels[index]

    def gaussian_blur_transform(self, image, radius):
        return pil2np(np2pil(image).filter(ImageFilter.GaussianBlur(radius))) / 255

    def image_compress_transform(self, image, quality):
        buf = BytesIO()
        np2pil(image).save(buf, format = 'jpeg', quality = quality)
        buf.seek(0)
        pil_image_jpeg = Image.open(buf).copy()
        buf.close()
        return pil2np(pil_image_jpeg)

    # Misclassification Ratio
    def MR(self):
        cnt = 0

        for i in range(len(self.adv_images)):
            if self.successful_adv(i):
                cnt += 1

        return cnt / len(self.adv_images)

    # Average Confidence of Adversarial Class
    def ACAC(self):
        cnt = 0
        conf = 0

        for i in range(len(self.adv_images)):
            if self.successful_adv(i):
                cnt += 1
                conf += np.max(self.softmax_outs[i])

        return conf / cnt

    # Average Confidence of True Class
    def ACTC(self):
        cnt = 0
        true_conf = 0

        for i in range(len(self.adv_images)):
            if self.successful_adv(i):
                cnt += 1
                true_conf += self.softmax_outs[i, self.labels[i]]

        return true_conf / cnt

    # Average L_p Distortion
    def ALD(self):
        ori_r = np.round(self.images * 255)
        adv_r = np.round(self.adv_images * 255)

        NUM_PIXEL = int(np.prod(self.images.shape[1:]))

        pert0 = self.images - self.adv_images
        pert = adv_r - ori_r

        dist_l0 = 0
        dist_l2 = 0
        dist_li = 0

        cnt = 0

        for i in range(len(self.adv_images)):
            if self.successful_adv(i):
                cnt += 1
                dist_l0 += np.linalg.norm(np.reshape(pert[i], -1), ord = 0) / NUM_PIXEL
                dist_l2 += np.linalg.norm(np.reshape(pert0[i], -1), ord = 2)
                dist_li += np.linalg.norm(np.reshape(pert0[i], -1), ord = np.inf)

        adv_l0 = dist_l0 / cnt
        adv_l2 = dist_l2 / cnt
        adv_li = dist_li / cnt

        return adv_l0, adv_l2, adv_li

    # Average Structural Similarity
    def ASS(self):
        ori_r_channel = np.transpose(np.round(self.images * 255), (0, 2, 3, 1))
        adv_r_channel = np.transpose(np.round(self.adv_images * 255), (0, 2, 3, 1))

        totalSSIM = 0
        cnt = 0

        for i in range(len(self.adv_images)):
            if self.successful_adv(i):
                cnt += 1
                totalSSIM += SSIM(ori_r_channel[i], adv_r_channel[i], multichannel = True)

        return totalSSIM / cnt

    # Perturbation Sensitivity Distance
    def PSD(self):
        psd = 0
        cnt = 0

        for outer in range(len(self.adv_images)):
            if self.successful_adv(outer):
                cnt += 1

                image = self.images[outer]
                pert = abs(self.adv_images[outer] - self.images[outer])

                for idx_channel in range(image.shape[0]):
                    image_channel = image[idx_channel]
                    pert_channel = pert[idx_channel]

                    image_channel = np.pad(image_channel, 1, 'reflect')
                    pert_channel = np.pad(pert_channel, 1, 'reflect')

                    for i in range(1, image_channel.shape[0] - 1):
                        for j in range(1, image_channel.shape[1] - 1):
                            psd += pert_channel[i, j] / np.std(np.array(
                                [image_channel[i - 1, j - 1], image_channel[i - 1, j], image_channel[i - 1, j + 1], image_channel[i, j - 1],
                                 image_channel[i, j], image_channel[i, j + 1], image_channel[i + 1, j - 1], image_channel[i + 1, j],
                                 image_channel[i + 1, j + 1]]))

        return psd / cnt

    # Noise Tolerance Estimation
    def NTE(self):
        nte = 0
        cnt = 0

        for i in range(len(self.adv_images)):
            if self.successful_adv(i):
                cnt += 1
                sort_preds = np.sort(self.softmax_outs[i])
                nte += sort_preds[-1] - sort_preds[-2]

        return nte / cnt

    # Robustness to Gaussian Blur
    def RGB(self):
        radius = 0.5

        total = 0
        num_gb = 0

        for i in range(len(self.adv_images)):
            if self.successful_adv(i):
                total += 1
                adv_sample = self.adv_images[i]
                gb_sample = self.gaussian_blur_transform(adv_sample, radius = radius)
                gb_pred = self.model.predict(np.array([gb_sample]))
                if self.attack.targeted:
                    if np.argmax(gb_pred) == self.targets[i]:
                        num_gb += 1
                else:
                    if np.argmax(gb_pred) != self.labels[i]:
                        num_gb += 1

        return num_gb / total

    # Robustness to Image Compression
    def RIC(self):
        quality = 90

        total = 0
        num_ic = 0

        for i in range(len(self.adv_images)):
            if self.successful_adv(i):
                total += 1
                adv_sample = self.adv_images[i]
                ic_sample = self.image_compress_transform(adv_sample, quality = quality)
                ic_pred = self.model.predict(np.array([ic_sample]))
                if self.attack.targeted:
                    if np.argmax(ic_pred) == self.targets[i]:
                        num_ic += 1
                else:
                    if np.argmax(ic_pred) != self.labels[i]:
                        num_ic += 1

        return num_ic / total

def main(model_name, attack_name, attack_config, suffix = ''):
    # 设置随机数种子
    set_seed()

    # 创建网络
    model = get_model(model_name)

    # 加载攻击
    attack = get_attack(attack_name, model, attack_config)

    # 加载评估
    ev = AttackEvaluation(model, attack, suffix = suffix)

    # 进行评估
    MR = ev.MR()
    ACAC = ev.ACAC()
    ACTC = ev.ACTC()
    ALD0, ALD2, ALDi = ev.ALD()
    ASS = ev.ASS()
    PSD = ev.PSD()
    NTE = ev.NTE()
    RGB = ev.RGB()
    RIC = ev.RIC()
    print('MR:\t{:.1f}%'.format(MR * 100))
    print('ACAC:\t{:.3f}'.format(ACAC))
    print('ACTC:\t{:.3f}'.format(ACTC))
    print('ALD0:\t{:.3f}'.format(ALD0))
    print('ALD2:\t{:.3f}'.format(ALD2))
    print('ALDi:\t{:.3f}'.format(ALDi))
    print('ASS:\t{:.3f}'.format(ASS))
    print('PSD:\t{:.3f}'.format(PSD))
    print('NTE:\t{:.3f}'.format(NTE))
    print('RGB:\t{:.3f}'.format(RGB))
    print('RIC:\t{:.3f}'.format(RIC))

    return {
        'MR': MR,
        'ACAC': ACAC,
        'ACTC': ACTC,
        'ALD0': ALD0,
        'ALD2': ALD2,
        'ALDi': ALDi,
        'ASS': ASS,
        'PSD': PSD,
        'NTE': NTE,
        'RGB': RGB,
        'RIC': RIC
    }
